'''
Author: kun 56216004@qq.com
Date: 2023-06-27 16:23:32
LastEditors: kun 56216004@qq.com
LastEditTime: 2023-06-27 16:50:10
FilePath: \langchain\ask1.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
# 专有模型，例如：OpenAI
# pip install openai
from langchain.llms import OpenAI
import openai
openai.api_base = "https://api.chatanywhere.com.cn/v1"
from config import * 
# openai.api_key = api_key
import os
os.environ["OPENAI_API_KEY"] = api_key # 填入OpenAI Secret Key

llm = OpenAI(model_name="text-davinci-003")


from langchain.document_loaders import TextLoader

loader = TextLoader('./data/产品经理.txt', encoding='utf8')

documents = loader.load()


from langchain.text_splitter import CharacterTextSplitter

text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

texts = text_splitter.split_documents(documents)


# 专有文本嵌入模型，例如：OpenAI
# pip install tiktoken
from langchain.embeddings import OpenAIEmbeddings
embeddings = OpenAIEmbeddings()

# pip install faiss-cpu
from langchain.vectorstores import FAISS

# 创建 vectorestore 用作索引
db = FAISS.from_documents(texts, embeddings)


from langchain.chains import RetrievalQA

retriever = db.as_retriever()

qa = RetrievalQA.from_chain_type(
    llm=llm, 
    chain_type="stuff", 
    retriever=retriever, 
    return_source_documents=True)

query = u"产品经理的工作责任是什么?"
result = qa({
   "query": query})

print(result['result'])